// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#pragma once
#include <memory>
#include <stack>
#include <string>
#include <utility>
#include <vector>

#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/asio/ssl/context.hpp>
#include <boost/asio/strand.hpp>
#include <boost/beast/core/tcp_stream.hpp>
#include <boost/beast/ssl/ssl_stream.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/websocket/stream_fwd.hpp>
#include <glog/stl_logging.h>

#include "coin2/base/log.h"
#include "coin2/base/rate_limiter.h"
#include "coin2/base/scheduler.h"
#include "coin2/exchange/base/feed/latency_extractor.h"
#include "coin2/flow/topic_string.h"
#include "coin2/exchange/base/feed/ws/feed_subscription_spec.h"

/// if a payload is compressed, return the decompressed string.
/// otherwise, return the original string w/ std::move.

// a list of raw strings to send to the exchange.

/// WS has control frames, including ping and pong.
/// boost::beast automatically handle these (e.g. ping frame is not visible to async_read).
/// however, some exchanges use the raw string 'ping' and 'pong' as payloads.
/// if IsSupportingPingFrame() is true,

/// if IsSupportingPingFrame() is false, handle ping yourself.
/// if not ping, return std::nullopt.

namespace coin2::exchange::feed {
class WebSocketConfig;
}
namespace coin2::exchange::base::feed {
class RequestGenerator;
}

namespace coin2::exchange::base::executor {
class LiveAsyncExecutor;
}

class AsyncLeakyBucket;

struct FeedSubscriptionSpec;

///////////////////////////////////////// ^^^ forward decls

namespace coin2::exchange::base::feed {
namespace impl {

using namespace std::string_literals;

class WsConnection : public std::enable_shared_from_this<WsConnection> {
 public:
  using WebsocketSsl = boost::beast::websocket::stream<
      boost::beast::ssl_stream<boost::beast::tcp_stream>>;

  using WebsocketNoSsl =
      boost::beast::websocket::stream<boost::beast::tcp_stream>;

  WsConnection(const WsConnectionConfig& config,
               std::function<void()> fail_callback,
               executor::LiveAsyncExecutor* exec,
               RequestGenerator* request_generator,
               std::shared_ptr<LatencyExtractor> latency_extractor,
               bool enable_telemetry);

  ~WsConnection();

  void RunAsync();

  std::string Decompress(const std::string& s) const;

  /// WS has control frames, including ping and pong.
  /// boost::beast automatically handle these (e.g. ping frame is not visible to
  /// async_read). however, some exchanges use the raw string 'ping' and 'pong'
  /// as payloads. if IsSupportingPingFrame() is true,
  std::optional<std::string> HandlePing(const std::string& s) const;
  std::vector<std::string> GetChannels() const;
  std::string GetSubscriptionRequest(const std::string& symbol_channel);
  std::vector<std::string> GetSubscriptionRequests() const;
  std::vector<std::string> GetUnsubscriptionRequests() const;
  void Unsubscribe(const std::vector<std::string>& symbols);
  void PostUnsubcribe(const std::vector<std::string>& symbols);
  std::vector<std::string> GetUnsubscriptionRequest(const std::vector<std::string> symbols) const;

  void fail(boost::asio::yield_context* yield, boost::system::error_code ec,
            const std::string& what);

  void fail_async1(boost::system::error_code ec, const std::string& what);
  void fail_async2(boost::system::error_code ec) { fail_callback_(); }

  /* --------------------------------------------------------------------------
   */
  /*                                    write */
  /* --------------------------------------------------------------------------
   */

  void Send(const std::string& msg);
  void SendImpl(boost::beast::error_code ec, const std::string& last_msg);

  void Subscribe(boost::asio::yield_context yield, const std::string& request);

  // send subscribe msgs in a spawned asio loop.
  // Since multiple subscriptions may take a long time,
  // we need to start reading before all subscriptions, otherwise the exchange
  // disconnects.
  void SubscribeAll(boost::asio::yield_context yield);

  void Unsubscribe() {
    std::weak_ptr<WsConnection> self(shared_from_this());
    boost::asio::spawn(strand_, [self](boost::asio::yield_context yield) {
      if (auto s = self.lock()) {
        s->UnsubscribeAll(std::move(yield));
      } else {
        LOG(ERROR)
            << "boost::asio handler object (WsConnection) does not exist";
      }
    });
  }

  void UnsubscribeAll(boost::asio::yield_context yield);

  bool IsUnsubscriptionResponse(const std::string& msg,
                                std::string* symbol_channel) const;

  void CleanUp();
  void PostCleanUp();

  int64_t Delay();

  virtual void Run(boost::asio::yield_context yield);

  void RunAfterConnect(boost::asio::yield_context yield, const std::string&,
                       const std::string&);

  void Run(boost::asio::yield_context yield, const std::string&,
           const std::string&, const std::string&);

 protected:
  virtual void SchedulePingPong();
  template <typename T>
  void CheckAndSetDecorator(T& stream);
  void InsertHeader(boost::beast::websocket::request_type& req);

 protected:
  WsConnectionConfig config_;
  std::function<void()> fail_callback_;
  volatile bool running_;
  volatile bool ws_closed_;
  executor::LiveAsyncExecutor* const exec_;

 private:
  bool IsSymbolRemoved(const std::string& symbol_channel) const;

  boost::asio::io_context::strand strand_;

  boost::asio::ssl::context ctx_;
  boost::asio::ip::tcp::resolver resolver_;

  bool use_ssl_;
  std::unique_ptr<WebsocketSsl> ws_ssl_;
  std::unique_ptr<WebsocketNoSsl> ws_nossl_;

  std::shared_ptr<AsyncLeakyBucket> unauthorized_subscription_limit_;
  std::shared_ptr<AsyncLeakyBucket> unauthorized_command_limit_;

  volatile bool sending_ = false;
  std::stack<std::string> send_stack_;

  // metric
  int64_t num_messages_ = 0;
  int64_t num_subscription_sent_ = 0;

 protected:
  std::shared_ptr<Scheduler> scheduler_;

 private:
  RequestGenerator* request_generator_;

  std::shared_ptr<LatencyExtractor> latency_extractor_;
  bool enable_telemetry_;

  std::vector<std::string> removed_symbols_;  // for unsubscribe response
};

}  // namespace impl

using impl::WsConnection;

}  // namespace coin2::exchange::base::feed
